# Wei Jiang

from twisted.internet.protocol import Factory
from twisted.internet.protocol import Protocol
from twisted.internet import reactor, protocol
from twisted.protocols.basic import LineReceiver
from twisted.internet.endpoints import TCP4ServerEndpoint
from optparse import OptionParser
from operator import itemgetter

import pickle
import random
import sys


class brokerServerCommand(LineReceiver):
    def __init__(self,factory,addr):
        self.actionCode = factory.actionCode
        self.brokerServerStateCode = None
        self.clientAddr = addr
        self.currentCommandCode = -1
        self.dataPortProtocolReference = None
        self.factory = factory
        
        if self.actionCode == 0:
            self.singleQueryID = 0
        elif self.actionCode == 1:
            self.channelQueriesDict = factory.channelQueriesDict
            self.queries = factory.queries
        
        self.clientConnections = factory.clientConnections
        self.commands = factory.commands
        self.channels = factory.channels
        self.brokerServerMachineStates = factory.brokerServerMachineStates
        self.clientConnectionsWithCorrespondingChannels = factory.clientConnectionsWithCorrespondingChannels
        self.command = None
        
        self.resultsMixer = factory.resultsMixer
        
        self.brokerServerStateCode = 0
        
        if self.actionCode == 0:
            self.resultsMixer[ self.singleQueryID ] = []
            pass
        elif self.actionCode == 1:
            for queryID in self.queries:
                if queryID not in self.resultsMixer:
                    self.resultsMixer[queryID] = []
                else:
                    # may have duplicate queryID cause this self.resultsMixer are shared by all the objects.
                    pass

    def test(self):
        print "I can control."
        
    def connectionMade(self):
        print "Get connection from:",self.clientAddr
        self.factory.numCommandProtocols = self.factory.numCommandProtocols + 1
        print "Currently",self.factory.numCommandProtocols,"slave machines online."
        self.brokerServerStateCode = 2
        
    def connectionLost(self, reason):
        if self.clientConnections.has_key(str(self.clientAddr)):
            del self.clientConnections[str(self.clientAddr)]
            del self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)]
            self.factory.numCommandProtocols = self.factory.numCommandProtocols - 1
        print "Currently",self.factory.numCommandProtocols,"slave machines online."
    
    def lineReceived(self, line):
        #It will receive different kinds of content. Will be (1) command code or (2)Actual content of the query
        #print "        ---->Received Line:",line
        try:
            #self.currentCommandCode = -1
            currentContent = ""
            if len(line.strip().split(" ") ) == 1:
                self.currentCommandCode = int(line.strip().split(" ")[0])
            elif len( line.strip().split(" ") ) == 2:
                self.currentCommandCode = int(line.strip().split(" ")[0])
                currentContent = line.strip().split(" ")[1]
            else:
                self.currentCommandCode = int(line.strip().split(" ")[0])
                for element in line.strip().split(" ")[1:]:
                    currentContent += element + " "
            
            currentContent = currentContent.strip()
            
            if self.currentCommandCode in self.commands:
                print "        ---->Received Command code:",self.currentCommandCode
                if self.currentCommandCode == 0:
                    print "        ---->Received Command content:",currentContent
                    self.handle_SubscribeChannels(currentContent)
                elif self.currentCommandCode == 1:
                    print "        ---->Received Command content:",currentContent
                    self.handle_UnsubscribeChannels(currentContent)        
                elif self.currentCommandCode == 2:
                    self.handle_GetStatus()
                elif self.currentCommandCode == 3:
                    self.handle_RegisterMyself()
                elif self.currentCommandCode == 4:
                    self.handle_GetChannels()
                elif self.currentCommandCode == 5:
                    self.handle_QueryMe()
                elif self.currentCommandCode == 6:
                    self.handle_QueryOthers()
                elif self.currentCommandCode == 7:
                    self.handle_IAmReady()
                elif self.currentCommandCode == 8:
                    print "        ---->Received Command content:",currentContent
                    self.handle_QueryOthersOneQueryAtATime(currentContent)
            else:
                print "Command code:",currentCommandCode,"NOT supported."
                self.sendLine("FAIL")
        except:
            #deal with the returning results.
            #print self.currentCommandCode
            print "Results Returning Section:"
            print "self.currentCommandCode:",self.currentCommandCode
            
            if self.currentCommandCode == 5:
                print "********************"
                queryResultDict = pickle.loads(line.strip())
                #print queryResultDict
                
                for channelNumber in queryResultDict:
                    print channelNumber,len(queryResultDict[channelNumber])
                
                '''
                for channelNumber in queryResultDict:
                    print channelNumber
                    for queryResultsTuple in queryResultDict[channelNumber]:
                        print queryResultsTuple
                '''
                print "********************"
            
            elif self.currentCommandCode == 7:
                print "********************"
                print "Add the results coming from:",str(self.clientAddr)
                queryResultDict = pickle.loads(line.strip())
                for channelNumber in queryResultDict:
                    #print "    channelNumber:",channelNumber
                    for (queryID,resultsList) in queryResultDict[channelNumber]:
                        #print "        Search queryID:",queryID,"with the query content:",self.queries[queryID]
                        #print "        Length of Results:",len(resultsList)
                        for (docID,score,url) in resultsList:
                            #print "            Score: ",score,"DocID: ",docID,"URL:",url
                            if (docID,score,url,self.clientAddr) not in self.resultsMixer[queryID]:
                                self.resultsMixer[queryID].append( (docID,score,url,self.clientAddr))
                            else:
                                print "The tuple we want to add (docID,score,url,self.clientAddr) has already existed in the dict."
                                print "The tuple is:","docID",docID,"score",score,"url",url,"self.clientAddr",self.clientAddr
                            
                print "Final results mix output:"
                for queryID in self.resultsMixer:
                    print queryID,len( self.resultsMixer[queryID] )
                    self.resultsMixer[queryID].sort(cmp=None, key=itemgetter(1), reverse=True)
                    for (docID,score,url,self.clientAddr) in self.resultsMixer[queryID]:
                        print "    ---->","Score:",score,"DocID:",docID,"URL:",url,"Source:",self.clientAddr
                print "********************"
            else:
                pass    
            
    
    def rawDataReceived(self,data):
        print "rawDataReceived:",data

    def handle_IAmReady(self):
        print "handle_IAMReady"
        if self.clientConnections[ str(self.clientAddr) ][1] == "UndefinedStatus":
            del self.clientConnections[ str(self.clientAddr) ][1]
            self.clientConnections[ str(self.clientAddr) ].append("Ready")
            self.sendLine("SUCCESS")
        else:
            self.sendLine("FAIL")
        
        
        print "        ---#>The broker actionCode is:",self.actionCode
        if self.actionCode == 0:
            print "        ---#>The broker is now waiting for some machines to active the interactive mode of distributed query processing(A little strange and will change in the future soon)..."
            
        elif self.actionCode == 1:
            print "        ---#>The broker is now waiting for some machines to active the batch mode of query evaluation(A little strange and will change in the future soon)..."
    
    
    def chooseChannelsBasedOnQuery(self,queryContent):
        channelNumberList = []
        #channelNumberList.append(0)
        #channelNumberList.append(1)
        #channelNumberList.append(3)
        
        
        #currently, just random choose the number and the list.
        numberOfChannels = random.randint(1, 4)
        #print "numberOfChannels:",numberOfChannels
        for i in range(0,numberOfChannels):
            #print "i:",i
            channelNumber = random.randint(0, 3)
            if channelNumber not in channelNumberList:
                channelNumberList.append( channelNumber )
        channelNumberList.sort(cmp=None, key=None, reverse=False)
        
        
        print "channelNumberList:",channelNumberList
        return channelNumberList
    
    def handle_QueryOthersOneQueryAtATime(self,queryContent):
        print "handle_QueryOthersOneQueryAtATime"
            
        if self.actionCode == 0:
            print "queryContent:",queryContent
            channelNumberListForThisQuery = self.chooseChannelsBasedOnQuery(queryContent)
            
            for endpoint in self.clientConnections:
                if self.clientConnections[endpoint][0] != self and self.clientConnections[endpoint][1] == "Ready":
                    print endpoint
                    #print self.clientConnections[endpoint]
                    outgoingQueryDict = {}
                    for channelNumber in self.clientConnectionsWithCorrespondingChannels[endpoint]:
                        if channelNumber in channelNumberListForThisQuery:
                            outgoingQueryNumberANDContentList = []
                            top_k = random.randint(1, 6 )
                            outgoingQueryNumberANDContentList.append( (self.singleQueryID,queryContent,top_k) )
                            
                            outgoingQueryDict[channelNumber] = outgoingQueryNumberANDContentList                
                    
                    
                    if len(outgoingQueryDict) != 0:
                        print "content to send:",outgoingQueryDict 
                        self.clientConnections[endpoint][0].sendLine(pickle.dumps( outgoingQueryDict ))
                    else:
                        print "content will NOT send to this slave machine."
                        print "reason:channels Not subscribed."
                        #print "outgoingQueryDict:",outgoingQueryDict
            
        elif self.actionCode == 1:
            print "Currently, there is nothing to do."
        
        #No matter how, send back the "SUCCESS" back to the command dispatcher. 
        self.sendLine("SUCCESS")      

    def handle_QueryOthers(self):
        print "handle_QueryOthers"
        #Let's assume that the slave machine know the broker's protocal very well.
        if self.actionCode == 0:
            print "Currently, there is nothing to do."
        elif self.actionCode == 1:
            for endpoint in self.clientConnections:
                if self.clientConnections[endpoint][0] != self and self.clientConnections[endpoint][1] == "Ready":
                    print endpoint
                    #print self.clientConnections[endpoint]
                    outgoingQueryDict = {}
                    for channelNumber in self.clientConnectionsWithCorrespondingChannels[endpoint]:
                        #print channelNumber
                        
                        outgoingQueryNumberANDContentList = []
                        
                        for queryNumber in self.channelQueriesDict[channelNumber]:
                            top_k = random.randint(1, 6 )
                            outgoingQueryNumberANDContentList.append( (queryNumber,self.queries[queryNumber],top_k) )
                        
                        outgoingQueryDict[channelNumber] = outgoingQueryNumberANDContentList                
                    
                    print "content to send:",outgoingQueryDict 
                    self.clientConnections[endpoint][0].sendLine(pickle.dumps( outgoingQueryDict ))
                
        

    def handle_QueryMe(self):
        print "handle_QueryMe:"
        outgoingQueryDict = {}
        self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)].sort()
        for channelNumber in self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)]:
            #print channelNumber
            outgoingQueryNumberANDContentList = []
            for queryNumber in self.channelQueriesDict[channelNumber]:
                outgoingQueryNumberANDContentList.append( (queryNumber,self.queries[queryNumber]) )
            outgoingQueryDict[channelNumber] = outgoingQueryNumberANDContentList
        self.sendLine( pickle.dumps(outgoingQueryDict) )
        
        
    def handle_SubscribeChannels(self,currentContent):
        channelCode = int(currentContent.strip())
        print "handle_SubscribeChannels:","select channel",channelCode,"...",
        if channelCode in self.channels and channelCode not in self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)]:
            self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)].append( channelCode )
            self.sendLine("SUCCESS")
            print "Done."
        else:
            print str(channelCode)," is NOT in the channelCode list OR have already subscribed.",
            self.sendLine("FAIL")
            print "NOT Done."
                
    def handle_UnsubscribeChannels(self,currentContent):
        channelCode = int(currentContent.strip())
        print "handle_UnsubscribeChannels:","delete channel",channelCode,"...",
        if channelCode in self.channels and channelCode in self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)]:
            self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)].remove( channelCode )
            self.sendLine("SUCCESS")
            print "Done."
        else:
            print str(channelCode)," is NOT in the list.",
            self.sendLine("FAIL")
            print "NOT Done."
           
    def handle_GetStatus(self):
        print "handle_GetStatus","...",
        self.sendLine( pickle.dumps(self.clientConnectionsWithCorrespondingChannels) )
        print "Done."
    
    def handle_GetChannels(self):
        print "handle_GetChannels"
        self.sendLine( pickle.dumps(self.channels) )

    def handle_DELCHANNEL(self,line):
        code = int( line.strip() )
        if code in self.clientConnectionsWithCorrespondingChannels[self.name]:
            print "Channel",code,"removed..."
            self.clientConnectionsWithCorrespondingChannels[self.name].remove( code )
        else:
            print str(code)," is NOT in the subscribed channel list."
            outputLine = "Select again from "
            print outputLine.strip(),str( self.clientConnectionsWithCorrespondingChannels[self.name] )
            
    def handle_RegisterMyself(self):
        print "handle_RegisterMyself:",self.clientAddr,"...",
        
        if self.clientConnections.has_key(str(self.clientAddr)):
            print "Connection Duplicated:",str(self.clientAddr)
            self.sendLine("Fail")
            print "Not Done."
        else:
            self.clientConnections[str(self.clientAddr)] = []
            self.clientConnections[str(self.clientAddr)].append(self)
            self.clientConnections[str(self.clientAddr)].append("UndefinedStatus")
            self.clientConnectionsWithCorrespondingChannels[str(self.clientAddr)] = []
            self.sendLine("SUCCESS")
            print "Done."

class brokerServerFactory(Factory):
    def __init__(self, typeCode, actionCode):
        self.actionCode = actionCode
        self.typeCode = typeCode
        self.numCommandProtocols = 0
        self.numDataProtocols = 0
        self.resultsMixer = {}
        
        if self.typeCode == 0:
            self.clientConnections = {} # maps user names to instances
            self.commands = {0:"SubscribeChannels", 1:"UnsubscribeChannels", 2:"GetStatus",3:"RegisterMyself",4:"GetChannels",5:"QueryMe",6:"QueryOthers",7:"IAMReady",8:"QueryOthersOneQueryPerTime"} # commands
            self.channels = {0:"News", 1:"Sports", 2:"Finance", 3:"Cars"} # channels
            self.brokerServerMachineStates = {0:"INIT",1:"END",2:"SERVE"}
            self.clientConnectionsWithCorrespondingChannels = {} # maps user names to their subscribe channels
            print "    ",len(self.commands),"commands have been set."
            print "    ",len(self.channels),"channels have been set."
            print "    ",len(self.brokerServerMachineStates),"machine states have been set."
            
            if self.actionCode == 0:
                #because it is trying to query in an interactive mode. There is nothing to do at this stage.
                pass
            elif self.actionCode == 1:
                #self.readQueriesFromDisk("/home/weijiang/workspace/polyIRToolkit/polyIRIndexer/1000query")
                self.readQueriesFromDisk("/data1/team/weijiang/machine-learning-project-related/programs/polyIRToolkit/polyIRIndexer/10query")
                self.queryAnalyzer()
            
        elif self.typeCode == 1:
            #currently, no use at all.
            pass

    def buildProtocol(self, addr):
        #print "incoming address:",addr
        if self.typeCode == 0:
            return brokerServerCommand(self,addr)
            
        elif self.typeCode == 1:
            #currently, no use at all.
            pass

    def queryAnalyzer(self):
        self.channelQueriesDict = {}
        #init
        for channel in self.channels:
            self.channelQueriesDict[channel] = []
        
        #assign the queries into different channels.
        #should be replaced by more advanced query analyzer.
        for queryNumber in self.queries:
            channelSelector = random.randint(0, len(self.channels)-1 )
            self.channelQueriesDict[channelSelector].append(queryNumber)
        
        print "    ",len(self.queries),"queries have been classified into",len(self.channels),"channels."
        for channelNumber in self.channelQueriesDict:
            print "        ",channelNumber,self.channels[channelNumber],len(self.channelQueriesDict[channelNumber]),self.channelQueriesDict[channelNumber]
     
    def readQueriesFromDisk(self,fileName):
        inputFileHandler = open(fileName,"r")
        self.queries = {}
        for index,line in enumerate( inputFileHandler.readlines() ):
            self.queries[index] = line.strip()
        print "    ",len(self.queries),"queries have been loaded."

def gotProtocol(p):
    #currently, no use.
    #print p.getHost()
    pass
    
        
    
    


############################################################################
#create an OptionParser instance
parser = OptionParser(version="%prog 2012/06/23")

#Typically, each option will have one short option string and one long option string
parser.add_option("-a", "--action", 
                  dest="broker_Action",
                  default="query",
                  type="string",
                  help="tell the broker to do perform what action, "
                        "Actions: query "
                        "or expert [default: %default]")

parser.add_option("--action-mode", 
                  dest="broker_Action_Mode",
                  default="interactive",
                  type="string",
                  help="tell the mode the broker is going to query the cluster machines, " 
                        "Interaction mode: interactive,batch, "
                        "or expert [default: %default]")

'''
parser.add_option("-q", "--quiet",
                  action="store_false", dest="verbose", default=True,
                  help="don't print status messages to stdout")
'''

(options, args) = parser.parse_args()

if options.broker_Action and options.broker_Action_Mode:
    print "The broker action is:",options.broker_Action
    print "The broker action mode is:",options.broker_Action_Mode
    if options.broker_Action == "query" and options.broker_Action_Mode == "interactive":
        actionCode = 0

############################################################################
#this broker factory can generate 2 kinds of products based on different type codes.
#typeCode: 0 for control port, 1 for data port
print "Program begins."

controlPort = 10877
controlPoint = TCP4ServerEndpoint(reactor,controlPort)

#typeCode in default is set to:0

#actionCode mapping table:
#0:query in an interactive mode.
#1:query in an batch mode.

typeCode = 0
actionCode = 0

deferredForCommandPort = controlPoint.listen(brokerServerFactory(typeCode,actionCode))
deferredForCommandPort.addCallback(gotProtocol)

print "Listening for incoming connections on port",str(controlPort)
print "Currently, 0 slave machines online."

reactor.run()
print "Program end."
